'''
Author: your name
Date: 2022-03-25 10:51:14
LastEditTime: 2022-03-26 22:49:27
LastEditors: Please set LastEditors
Description: 
'''
from collections import namedtuple
import akshare as ak
import pandas as pd
from typing import (List, Tuple, Union, Dict)


def get_topn2cons(cons_name: str,
                  rank_col: str,
                  concept_dic: Dict,
                  price: pd.DataFrame,
                  n=5,
                  cols: List = [1, 5, 10, 20, 60, 120]) -> pd.DataFrame:
    """获取指定概念 N日动量个股情况

    Parameters
    ----------
    cons_name : str
        指定概念
    rank_col : str
        按N日排名
    concept_dic : Dict
        概念字典
    price : pd.DataFrame
        价格数据

    Returns
    -------
    pd.DataFrame
    """
    target = concept_dic[0].get(cons_name, '')
    if target != '':
        target = [concept_dic[1][i] for i in target]

        cons_pct = calc_pct_nd(price, cols)
        return cons_pct.loc[target, :].sort_values(rank_col,
                                                   ascending=False).iloc[:n]


def get_ths_hist(ths_names: Union[List, str],
                 years: Union[List, str]) -> pd.DataFrame:
    """通过概念名称获取 概念数据

    Parameters
    ----------
    ths_names : Union[List,str]
        概念板块
    years : Union[List,str]
        _description_

    Returns
    -------
    pd.DataFrame
        _description_
    """
    if isinstance(ths_names, str):

        symbths_namesols = [ths_names]

    if isinstance(years, str):

        years = [years]

    dfs = []
    for year in years:

        for ths_name in ths_names:

            try:
                stock_board_concept_hist_ths_df = ak.stock_board_concept_hist_ths(
                    start_year=year, symbol=ths_name)
            except (AttributeError, KeyError, ValueError):
                print(ths_name)

            stock_board_concept_hist_ths_df['概念名称'] = ths_name

            dfs.append(stock_board_concept_hist_ths_df)

    data = pd.concat(dfs)

    return data


def calc_rs_gn(df: pd.DataFrame, N: int) -> pd.DataFrame:
    """计算相对强弱指标 RS

    Args:
        df (pd.DataFrame): index-date columns-symbol value
        N (int): 窗口期

    Raises:
        ValueError: 窗口不能大于df

    Returns:
        pd.DataFrame: N日RS
    """
    if len(df) <= N:
        raise ValueError('参数N不能大于df的长度')

    pct_chg = df.pct_change(N)  # N日收益率
    # 横向排序 降序 归一化
    rank = pct_chg.rank(axis=1, pct=True, ascending=True)

    return rank


def calc_rs(df: pd.DataFrame, col: str) -> pd.DataFrame:
    """计算相对强弱指标 RS

    Args:
        df (pd.DataFrame): index-date columns-symbol value
        col (int): 窗口期

    Raises:
        ValueError: 窗口不能大于df

    Returns:
        pd.DataFrame: N日RS
    """
    if col not in df.columns:
        raise ValueError('col不在df中')

    pct_chg = df[col]  # N日收益率
    # 横向排序 降序 归一化
    rank = pct_chg.rank(pct=True, ascending=True)

    return rank


def get_gn_concept(names: Union[str, List]) -> namedtuple:
    """获取概念成分股

    Returns
    -------
    namedtuple
        codes k为板块名称 v-codes
        sec_name k-codes v-sec_name
    """
    res = namedtuple('Res', 'codes,sec_names')
    if isinstance(names, str):
        names = [names]

    code_dic = {}
    sec_name = {}

    for name in names:

        try:
            df = ak.stock_board_concept_cons_ths(symbol=name)
        except (AttributeError, IndexError):
            print(name)
        dic = df.set_index('代码')['名称'].to_dict()
        dic = {
            k + '.SH' if k[0] == '6' else k + '.SZ': v
            for k, v in dic.items()
        }
        sec_name.update(dic)
        code_dic[name] = list(dic.keys())

    return res(codes=code_dic, sec_names=sec_name)


def calc_pct_nd(price: pd.DataFrame, nd: Union[int, List]) -> pd.DataFrame:
    """计算n_days的收益率情况

    Parameters
    ----------
    price : pd.DataFrame
        _description_
    nd : Union[int,List]
        _description_

    Returns
    -------
    pd.DataFrame
        _description_
    """
    if isinstance(nd, int):

        nd = [nd]

    cols = [f"{i}日" for i in nd]

    pct_chg: pd.DataFrame = pd.concat(
        (price.pct_change(i).iloc[-1] for i in nd), axis=1)
    pct_chg.columns = cols

    return pct_chg
